import torch
import torch.nn as nn
import torch.nn.functional as F

class BaseExpert(nn.Module):
    def __init__(self, d_inputs, d_output):
        super(self).__init__()
        self.fc = nn.Linear(d_inputs, d_output)

    def forward(self, x):
        x = self.fc(x)
        return x
    
class BaseMoE(nn.Module):
    def __init__(self, d_inputs, d_output, n_experts):
        super(BaseMoE, self).__init__()
        self.gate = nn.Linear(d_inputs, n_experts)
        self.experts = nn.ModuleList([BaseExpert(d_inputs, d_output) for _ in range(n_experts)])

    def forward(self, x):
        expert_weights = F.softmax(self.gate(x), dim=1)
        expert_out_list = [
            expert(x) for expert in self.experts
        ]
        export_output = torch.concat(
            [
                expert_out.unsqueeze(1) for expert_out in expert_out_list
            ],  dim=1
        )
        expert_weights = expert_weights.unsqueeze(1)
        output = (export_output @ expert_weights).squeeze(1)
        return output

